import numpy as np
import cv2 as cv
import matplotlib.pyplot as plt

img1 = cv.imread('/Users/wanggh/Desktop/icon1.png', cv.IMREAD_GRAYSCALE)  # queryImage
img2 = cv.imread('/Users/wanggh/Desktop/original.jpeg', cv.IMREAD_GRAYSCALE)  # trainImage
# 初始化SIFT描述符
sift = cv.SIFT_create()
# 基于SIFT找到关键点和描述符
kp1, des1 = sift.detectAndCompute(img1, None)
kp2, des2 = sift.detectAndCompute(img2, None)
# FLANN的参数
FLANN_INDEX_KDTREE = 1
index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
search_params = dict(checks=50)  # 或传递一个空字典
flann = cv.FlannBasedMatcher(index_params, search_params)
matches = flann.knnMatch(des1, des2, k=2)
# 只需要绘制好匹配项，因此创建一个掩码
matchesMask = [[0, 0] for i in range(len(matches))]
# 根据Lowe的论文进行比例测试
for i, (m, n) in enumerate(matches):
    if m.distance < 0.7 * n.distance:
        matchesMask[i] = [1, 0]
draw_params = dict(matchColor=(0, 255, 0),
                   singlePointColor=(255, 0, 0),
                   matchesMask=matchesMask,
                   flags=cv.DrawMatchesFlags_DEFAULT)
img3 = cv.drawMatchesKnn(img1, kp1, img2, kp2, matches, None, **draw_params)
cv.imshow("test", img3)
cv.waitKey(0)
cv.destroyAllWindows()
